# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question2/my_special.py
import math

import numpy
import scipy.constants
import sympy
import sympy.abc


# This function is copied from
# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question1/main.py
def _associated_laguerre(k, p_times_2, caches) -> sympy.Expr | int:
    key = (k, p_times_2)
    if key in caches:
        return caches[key]
    t = sympy.abc.t
    if k == 0:
        result = 1
    else:
        previous = _associated_laguerre(k - 1, p_times_2, caches)
        p = p_times_2 / 2
        result = sympy.diff(previous, t) * t
        result += (k + p - t) * previous
        result /= k
    caches[key] = result
    return result


_assoc_laguerre_caches = {}


def assoc_laguerre(x, n, k):
    expression = _associated_laguerre(n, k * 2, caches=_assoc_laguerre_caches)
    return sympy.lambdify(sympy.abc.t, expression)(x)


def sph_harm(m, n, theta, phi):
    # TODO: 阶乘或可优化
    # 是缓存阶乘结果进行下次运算更好，还是直接利用快速阶乘计算更快？
    result = (2 * n + 1) / 4 / numpy.pi
    result *= math.factorial(n - m) / math.factorial(n + m)
    result **= 1 / 2
    result *= numpy.exp(1j * m * theta)
    # 这边的连带勒让德函数还是直接用了 SciPy ，
    # 如果自己写的话大概率就是根据公式去算，
    # 而里面求导之类的也是肯定得由 SymPy 来完成，
    # 这边不想再额外引入符号运算系统了。
    result *= scipy.special.lpmv(m, n, numpy.cos(phi))
    return result
